package com.feidee.fd.sml.algorithm.ml

import com.feidee.fd.sml.algorithm.component.ml.regression.{IsotonicRegressionComponent, IsotonicRegressionParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.scalatest.FunSuite

/**
  * @Author: dongguosheng
  * @Date: 2019/3/25 17:36
  * @Review
  * @Email: guosheng_dong@sui.com
  */
class IsotonicRegressionComponentSuite extends FunSuite {

  val isotonic = new IsotonicRegressionComponent()
  val paramStr: String =
    """
      |{
      |	'input_pt': '',
      |	'output_pt': '',
      |	'featuresCol': 'features',
      |	'labelCol': 'label',
      | 'probabilityCol': 'probability',
      |	'modelPath': '',
      |	'metrics': ['accuracy', 'auc', 'pr', 'abc'],
      |	'isotonic': false,
      |	'featureIndex': 0
      |}
    """.stripMargin

  val param: IsotonicRegressionParam = isotonic.parseParam(new ToolClass().encrypt(paramStr))

  /**
    * IsotonicRegression 参数构造方法测试：检测到读取出合法参数 & 成功赋值默认参数，即认为测试通过
    */
  test("IsotonicRegression parameter construction test") {
    assert("".equals(param.input_pt) & "".equals(param.output_pt) & "features".equals(param.featuresCol) &
      "label".equals(param.labelCol) & "prediction".equals(param.predictionCol) &
      param.weightCol == null & !param.isotonic & "".equals(param.modelPath) &
      param.metrics.sameElements(Array("accuracy", "auc", "pr", "abc")) & param.featureIndex == 0
    )
  }

  /**
    * IsotonicRegression 模型训练功能测试：传入生成的测试数据和正确参数，可以正常训练 IsotonicRegression 模型，且预测结果列等于生成数据列数，即认为测试通过
    */
  test("IsotonicRegression training function test") {
    // 测试模型训练方法
    val testingData = TestingDataGenerator.sampleIsotonicRegressionData
    val model = isotonic.train(param, testingData)
    model.transform(testingData).show()
  }

}
